import os
import psutil
import requests
import argparse
from flask import Flask, request, jsonify
from sentence_transformers import SentenceTransformer
from datetime import datetime

app = Flask(__name__)

print("emb_server 开始初始化模型")
embModel = None
embModelPath = "E:\\models\\embedding\\bge-large-zh-v1.5"
if os.path.exists(embModelPath):
    embModel = SentenceTransformer(embModelPath)
if embModel:
    print("emb_server 初始化模型成功")
else:
    print("emb_server 初始化模型失败")
    raise Exception("emb_server 初始化模型失败")

@app.route('/', methods=['GET'])
def index():
    return "index"
@app.route('/v1/embedding', methods=['POST'])
def v1_embedding():
    ret = False
    msg = "未知错误"
    embedding = []

    params = request.get_json()
    print("[%s]v1_embedding() params=%s"%(str(datetime.now()),str(params)))
    try:
        content = params["content"]
        contents = [content]
        embeddings = embModel.encode(contents)
        embedding = embeddings[0]
        embedding = list(map(lambda x: float(x), embedding))

        # print("v1_embedding() embedding:", type(embedding), len(embedding), embedding)
        ret = True
        msg = "success"
    except Exception as e:
        msg = str(e)

    response_data = {
        "code": 1000 if ret else 0,
        "msg": msg,
        "embedding": embedding
    }
    return jsonify(response_data)


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Embedding Server')
    parser.add_argument('--host', default='0.0.0.0', help='服务器地址')
    parser.add_argument('--port', type=int, default=1235, help='服务器端口')
    parser.add_argument('--debug', type=int, default=0, help='运行模式')
    args = parser.parse_args()
    app.run(
        host=args.host,
        port=args.port,
        debug=args.debug,
        threaded=True
    )